"""
Created on 03-Aug-2009

@author: dexen
"""

import cPickle
from dexen.system.base.individual import Individual
from dexen.system.node.master.archive_mgr import ArchiveMgr
from dexen.system.base.spec import DexenSpec


class IndMgr(object):
    """A class to manage partially evaluated individuals in memory.
    
    """
    def __init__(self, spec):
        """Constructor
        
        Args:
        inds_db : A dict that stores all of the partially/fully evaluated
        individuals. Each individuals is made of dict as well. All the 
        individuals are pickled before stored and unpickled before retrieved.

        inds_db_lock : A shared lock for synchronizing access to the database.
    
        tasks : A dict that stores all the information about tasks.
        
        """
        assert isinstance(spec, DexenSpec)
        self.spec = spec
        self.inds_db = {}
        self.arc_mgr = ArchiveMgr(self.spec)
        self.tasks = spec.tasks
        self.auto_id = 1 # WARNING: ! auto_id is supposed to be synchronized
        self._n_live_inds = 0
    
    def get_ind(self, id):
        ind = cPickle.loads(self.inds_db[id])
        return ind
        
    def _isvalid_ind(self, ind, task):
        if not ind.__isalive__():
            return False 
        obj = task.get_obj()
        return obj.__isvalid__(ind)
    
    def _determine_queues(self, ids, ind):
        """Assign the id of ind to the queues that can process.
        
        Args:
        ids: A dict to store the list of ids for each task queue. 
        
        ind: An individual to be tested against all the task input format.
        
        Note:
        ids will be returned by set_inds to queue manager. Then the 
        queue manager will put the ids into queues accordingly.
        
        """
        for task_name, task in self.tasks.items():
            if self._isvalid_ind(ind, task):
                ids[task_name].append(ind.__id__)

    def _insert_ind(self, ind):
        assert isinstance(ind, Individual)
        ind.__id__ = self.auto_id #TODO: check if added to __modified__
        ind.revive() #TODO: check
        ind.__modified__ = {}
        
        #print "In _insert_ind, the birth id", ind.__birth_id__
        
        self.arc_mgr.archive_data(ind)
        ind.strip_archived()
        
        self.inds_db[self.auto_id] = cPickle.dumps(ind)
        self.auto_id += 1
        
    def _update_ind(self, ind):
        target_ind = cPickle.loads(self.inds_db[ind.__id__])
        assert isinstance(target_ind, Individual)

        for key in ind.__modified__.keys():
            #print "modified attribute:", key, "\n"
            setattr(target_ind, key, getattr(ind, key)) # TODO: check if works
            #print "The modified, ", key
        target_ind.__modified__ = {} #WARNING: important to reset    
        
        self.arc_mgr.archive_data(target_ind)
        target_ind.strip_archived()
        
        #if not target_ind.__isalive__():
        #    target_ind.strip_all()
        
        self.inds_db[ind.__id__] = cPickle.dumps(target_ind)
        return target_ind

    def get_inds(self, task_name, ids):
        """Get individuals out of the inds_db database.
        
        task_name is used to select the corresponding fields of the 
        individuals.
            
        """
        inds = []
        for id in ids:
            ind = cPickle.loads(self.inds_db[id])
            inds.append(ind)
        return inds

    def get_all_inds(self):
        inds = []
        for id, pickled_ind in self.inds_db.items():
            inds.append(cPickle.loads(pickled_ind))
        return inds

    def set_inds(self, inds):
        """Put resulting individuals back to the inds_db.
        
        Returns a dict 'res' composed of (iqueue_name, list of ids) pairs.
        
        The calculation of n live inds is based on the ASSUMPTION that the 
        dead ind will never be put back into the task input queues.
        
        """
        
        cur_id = self.n_cur_births()
        ids = dict.fromkeys(self.tasks, None)
        for task_name in ids.keys(): 
            ids[task_name] = []
        for ind in inds:
            if not ind._isnew():
                if not ind.__isalive__():
                    ind.__death_id__ = cur_id
                updated_ind = self._update_ind(ind)
                #print "Updated ind info\n", updated_ind
                if not updated_ind.__isalive__():
                    self._n_live_inds -= 1
                self._determine_queues(ids, updated_ind)
            else: 
                self._n_live_inds += 1
                ind.__birth_id__ = cur_id
                self._insert_ind(ind)
                self._determine_queues(ids, ind)
        return ids
    
    def n_cur_births(self):
        return self.auto_id-1
    
    def n_live_inds(self):
        return self._n_live_inds
    
    def __repr__(self):
        res = ""
        for id, ind in self.inds_db.items():
            ind = cPickle.loads(ind)
            res += "%s: %s\n"%(id, ind)
        return res
